import os
import random
from PIL import Image
from torch.utils.data import Dataset

random.seed(42)
# label_name = img_name.split('.')[0]
DogCat_label = {"Dog":0, "Cat":1}

class DogCatDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        """
        rmb面额分类任务的Dataset
        :param data_dir: str, 数据集所在路径
        :param transform: torch.transform，数据预处理
        """
        # data_info存储所有图片路径和标签，在DataLoader中通过index读取样本
        self.data_info = self.get_img_info(data_dir)
        self.transform = transform
    
    def __getitem__(self, index):
        # 通过 index 读取样本
        path_img, label = self.data_info[index]
        # 注意这里需要 convert('RGB')
        '''
        Image.open: https://blog.csdn.net/weixin_42213622/article/details/109110140
        convert https://blog.csdn.net/nienelong3319/article/details/105458742
        '''
        img = Image.open(path_img).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)

        return img, label
    
    def __len__(self):
        return len(self.data_info)

    '''
    使用@staticmethod修饰的方法是静态方法， 如果在方法中不需要访问任何实例方法和属性，纯粹地通过传入参数并返回数据的功能性方法，
    那么它就适合用静态方法来定义，它节省了实例化对象的开销成本，往往这种方法放在类外面的模块层作为一个函数存在也是没问题的，而放在类中，仅为这个类服务。
    例如下面get_img_info，它没有引用任何类或者实例相关的属性和方法。https://zhuanlan.zhihu.com/p/28010894
    关于 @classmethod详细使用场景 https://blog.csdn.net/qq_42217711/article/details/118147508?ydreferer=aHR0cHM6Ly9jbi5iaW5nLmNvbS8%3D
    '''
    @staticmethod
    def get_img_info(data_dir):
        data_info = []
        '''
        os.walk:用来遍历文件夹及子文件夹下所有文件并得到路径 https://blog.csdn.net/qq_41562433/article/details/82995098
        os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表。
        filter() 函数用于过滤序列，过滤掉不符合条件的元素，返回一个迭代器对象，如果要转换为列表，可以使用 list() 来转换。
        '''
        for root, dirs, _ in os.walk(data_dir):
            for sub_dir in dirs:
                img_names = os.listdir(os.path.join(root, sub_dir))
                img_names = list(filter(lambda x: x.endswith('.jpg'), img_names))
                for i in range(len(img_names)):
                    img_name = img_names[i]
                    path_img = os.path.join(root, sub_dir, img_name)
                    # label_name = img_name.split('.')[0]
                    label = DogCat_label[sub_dir]
                    data_info.append((path_img, int(label)))
        return data_info